# -*- coding: utf-8 -*-
"""
Created on Sat Dec  5 15:56:17 2020

@author: Team317
"""

import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
from resnet100 import KitModel
import PIL.Image as Image
import numpy as np
from face_modules.mtcnn import *
import cv2
import matplotlib.pyplot as plt
import pickle
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

#定义迭代器
trainloader = torch.utils.data.DataLoader(dataset = train_data,
                                          batch_size = 32,
                                          shuffle = True)
testloader = torch.utils.data.DataLoader(dataset = test_data,
                                          batch_size = 32,
                                          shuffle = True)



net = AgeNet()

#定义Loss和优化方法
#优化器和loss
optimizer = torch.optim.SGD(net.parameters(),lr = 0.001,momentum = 0.9)#优化器
loss_func = nn.CrossEntropyLoss()#计算损失函数


#训练开始
print("Start Training...")
for epoch in range(2):
    # 我们用一个变量来记录每100个batch的平均loss
    loss100 = 0.0
    # 我们的dataloader派上了用场
    for step, (batch_data,batch_label) in enumerate(trainloader):
        batch_data = Variable(batch_data)
        batch_label = Variable(batch_label)
        
        pre_label = net(batch_data)
        loss = loss_func(pre_label, batch_label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step() #更新参数
        loss100 += loss.item() #tensor.item()得到的是一个值
        #每隔100步先试一下训练效果
        if step % 100 == 0:
            print('[Epoch %d, Batch %5d] loss: %.3f' %
                  (epoch + 1, step + 1, loss100 / 100))
            loss100 = 0.0

print("Done Training!")

#保存神经网络
torch.save(net,'net.pkl')

